# # 参考 https://blog.csdn.net/greatsam/article/details/104260033
# # https://blog.csdn.net/jameschen9051/article/details/122106152
import copy
import os
import random
import cv2
import numpy as np


class Cv2Poly:
    polys = {}  # 会有多个多边形
    p_index = 0  # 当前索引
    poly_points = []  # 当前多边形
    drawing = True
    img = None  # 当前背景
    img_backup = None  # 重置的时候后使用
    file = None  # 背景图片文件路径

    @staticmethod
    def rand():
        return random.randrange(0, 255)

    def __init__(self):

        cv2.namedWindow('video')
        cv2.setMouseCallback('video', self._draw_ROI)  # 重点
        self.colors = [(self.rand(), self.rand(), self.rand(), .5) for _ in range(10)]

    def _draw_ROI(self, event, x, y, flags, param):  #
        if event == cv2.EVENT_LBUTTONDOWN:  #
            self.drawing = True
            if self.drawing:
                self.point1 = (x, y)
                self.poly_points.append((x, y))  # 存起来用于画多边形
                cv2.circle(self.img, self.point1, 5, (0, 255, 0), 2)  #
                self.refresh_canvas()  # 刷新画布
        if event == cv2.EVENT_RBUTTONDOWN:
            if self.drawing:
                self.pts = np.array([self.poly_points], np.int32)
                self.draw_poly(self.img, self.pts, self.p_index)
                self.refresh_canvas()
                self.finish_current_poly()

    def draw_poly(self, bk, pts, index):
        color = self.colors[index % len(self.colors)] #循环使用颜色
        cv2.fillPoly(bk, pts, color=color)
        self.drawing = False  # 基本没用

    def finish_current_poly(self):
        self.next_poly()

    def next_poly(self):
        self.polys[str(self.p_index)] = [i for i in self.poly_points]
        self.p_index += 1
        self.poly_points = []

    def draw_poly_on_black(self):  # 画透明背景多边形
        bk = np.zeros(self.img.shape)
        polys = self.polys.values()
        for i, p in enumerate(polys):
            ps = np.array([p], np.int32)
            self.draw_poly(bk, ps, i)
        return bk

    def save_file(self, bk):  #保存多边形

        # folder = os.path.dirname(self.file) if self.file is not None else os.path.dirname(__file__)
        folder = os.path.dirname(self.file)
        stem = os.path.basename(self.file).split('.')[0]
        to_file = f'{folder}/{stem}_mask.png'
        print("保存",to_file)
        cv2.imwrite(to_file, bk)

    def draw_all_poly(self):  # 重画多边形
        self.img = copy.deepcopy(self.img_backup)
        polys = self.polys.values()
        for i, p in enumerate(polys):
            # ps= [[*ip] for ip in p]
            ps = np.array([p], np.int32)
            self.draw_poly(self.img, ps, i)
        self.refresh_canvas()

    def refresh_canvas(self):
        cv2.imshow('video', self.img)

    def remove_last_poly(self):  #删除最后一个
        self.drawing = False
        if self.p_index == 0:
            return
        del (self.polys[str(self.p_index - 1)])
        self.p_index -= 1
        self.draw_all_poly()
        self.drawing = True

    def draw(self, img):
        if isinstance(img, str):
            self.img = cv2.imread(img)
            self.file = img
        else:
            self.img = img
            self.file = __file__  # 主要是用来取文件夹路径
        self.img_backup = copy.deepcopy(self.img)
        cv2.imshow('video', self.img)
        while True:
            # time.sleep(1)
            if cv2.waitKey(0) & 0xFF == ord('-'):  # 按-键删除
                self.remove_last_poly()
            if cv2.waitKey(0) & 0xFF == ord('s'):  # 按s键退出
                mask = self.draw_poly_on_black()

                self.save_file(mask)
            if cv2.waitKey(0) & 0xFF == ord('q'):  # 按q键退出
                break

if __name__ == '__main__':
    file = '/home/cjc/Pictures/girl.jpg'
    drawer = Cv2Poly()
    drawer.draw(file)